#!/usr/bin/env python 
# -*- coding:utf-8 -*-

import os,getopt,sys
from re import search,I
import time
from cmd import Cmd
from lib.utils.printer import *
import localapi
import lib.utils.settings as settings
from lib.utils.check import *
import traceback

class Cli(Cmd):
    def __init__(self):
        os.system("cls")
        Cmd.__init__(self)
        banner = '''
         _       __     __   _____
        | |     / /__  / /_ / ___/_________ _____
        | | /| / / _ \/ __ \\__ \/ ___/ __ `/ __  \\
        | |/ |/ /  __/ /_/ /__/ / /__/ /_/ / / / /
        |__/|__/\___/_.___/____/\___/\__,_/_/ /_/{}.{}
        '''.format(" by kali","p3.v.0.0.1")
        self.intro = banner
        self.prompt = 'webscan>> '
    def do_help(self,param):
        param = param.strip()
        if param == "set":
            localapi.help_set()
        else:
            localapi.info_plugins()
    def do_info(self,param:str):
        param = param.strip()
        if param in settings.plugs_info.keys():
            if param == "attacks":
                localapi.info_attacks()
            elif param == "audit":
                localapi.info_audit()
            elif param == "brute":
                localapi.info_brute()
            elif param == "disclosure":
                localapi.info_disclosure()
        else:
            warn("命令输入错误!请输入help查看!")
    def do_set(self,param):
        param = param.strip()
        if param.find(" ") > 0:
            params = param.split(" ")
            if len(params) ==2:
                key,value = param.split(" ")
                if key == "headers":
                    if value.find(":") > 0:
                        localapi.set_headers(value)
                    else:
                        warn("格式错误!")
                elif key == "auth":
                    localapi.set_auth(value)
                elif key == "cookie":
                    localapi.set_cookie(value)
                elif key == "timeout":
                    localapi.set_timeout(value)
                elif key == "proxy":
                    localapi.set_proxy(value)
                elif key == "pauth":
                    localapi.set_proxy_auth(value)
                elif key == "redirect":
                    localapi.set_redirect_false(value)
                elif key == "url":
                    url = Curl(value)
                    if search(settings.REG_URL,url):
                        localapi.set_url(url)
                    else:
                        warn("格式错误!")
                else:
                    warn("命令错误!")
            else:
                warn("命令错误!")
        
    def do_check(self,param):
        param = param.strip()
        if param == 'argv':
            localapi.get_argv()

    def do_exec(self,param:str):
        param  = param.strip()
        if not localapi.is_set_argv("url"):
            warn("URL设置是必选项!输入help set查看!")
        else:
            if param.find(".") > 0:
                params = param.split(".")
                if len(params) == 2:
                    plugin,mod = param.split(".")
                    if plugin in settings.plugs_info.keys():
                        if plugin == "attacks" and mod in settings.attacks_info.keys():
                            localapi.startup_spec_attacks(mod)
                        elif plugin == "audit" and mod in settings.audit_info.keys():
                            localapi.startup_spec_audit(mod)
                        elif plugin == "brute" and mod in settings.brute_info.keys():
                            localapi.startup_spec_brute(mod)
                        elif plugin == "disclosure" and mod in settings.disclosure_info.keys():
                            localapi.startup_spec_disclosure(mod)
                        else:
                            warn("无可用模块!")
                    else:
                        warn('无可用模块!')
            else:
                plugin = param
                if plugin in settings.plugs_info.keys():
                    if plugin == "attacks":
                        localapi.startup_full_attacks()
                    elif plugin == "audit":
                        localapi.startup_full_audit()
                    elif plugin == "brute":
                        localapi.startup_full_brute()
                    elif plugin == "disclosure":
                        localapi.startup_full_disclosure()
                    else:
                        warn("无可用模块!")
                else:
                    warn('无可用模块!')
    def do_report(self,param):
        filename = param.strip()
        if filename:
            localapi.generate_report(filename)
    def do_exit(self,param):
        warn("退出!")
        sys.exit()
    def default(self,param):
        print("命令不存在!")

if __name__ == "__main__":
    try:
        cli = Cli()
        cli.cmdloop()
    except Exception as e:
        #warn(e)
        print(e)
        traceback.print_exc()